package com.dyava.commons.token_bucket.local;

import com.dyava.commons.token_bucket.TokenBucketDefined;
import com.dyava.commons.token_bucket.TokenBucketException;
import com.dyava.commons.token_bucket.TokenBucketPool;

import java.util.concurrent.ConcurrentHashMap;

public class LocalTokenBucketPool extends TokenBucketPool {

    ConcurrentHashMap<String,LocalTokenBucket> bucketMap = new ConcurrentHashMap<>();

    @Override
    public void addDefined(TokenBucketDefined tokenBucketDefined){
        LocalTokenBucket tokenBucket = bucketMap.get(tokenBucketDefined.getName());
        if(tokenBucket==null){
            tokenBucket = new LocalTokenBucket(tokenBucketDefined.getName(),tokenBucketDefined.getLimitNum(),tokenBucketDefined.getLimitDen());
            LocalTokenBucket old = bucketMap.putIfAbsent(tokenBucketDefined.getName(),tokenBucket);
            if(old!=null){
                addDefined(tokenBucketDefined);
            }
        }
    }

    @Override
    public LocalTokenBucket getBucket(String name) {
        LocalTokenBucket localTokenBucket = bucketMap.get(name);
        return localTokenBucket;
    }

    @Override
    public boolean limit(String name, Object id) {
        LocalTokenBucket localTokenBucket = bucketMap.get(name);
        if(localTokenBucket==null){
            throw new TokenBucketException(name+"令牌桶不存在");
        }
        return localTokenBucket.limit(id);
    }
}
